import os
from os.path import join
import numpy as np
from visualize_egocentric_data import get_mocap_data
from utils.bvh2joint import bvh2joint, default_end_link_trans
from utils.fbx2objpose import fbx2objpose, fbx2camerapose
from utils.denoising import denoising_pose
from utils.visualization import render_mocap_data_on_camera_space
from utils.txt2intrinsic import txt2intrinsic
import cv2
import open3d as o3d


if __name__ == "__main__":

    cfg = {}

    ############## CHANGE BELOW #################
    # mocap data
    cfg["person1_bvh_path"] = "/home/liuyun/HHO-dataset/data_processing/exp_data/20230522_data/24/person1.bvh"
    cfg["person2_bvh_path"] = "/home/liuyun/HHO-dataset/data_processing/exp_data/20230522_data/24/person2.bvh"
    cfg["obj_fbx_path"] = "/home/liuyun/HHO-dataset/data_processing/exp_data/20230522_data/24/desk001.fbx"
    # camera parameters
    cfg["3rd2_intrinsic_path"] = "/home/liuyun/HHO-dataset/data_processing/camera_info/d435i_2/intrinsic.txt"
    cfg["3rd2_camera2world_path"] = "/home/liuyun/HHO-dataset/data_processing/camera_info/d435i_2/camera2world.txt"
    # object model
    cfg["obj_model_path"] = "/home/liuyun/HHO-dataset/data_processing/exp_data/objects/table001/desk1_m.obj"
    # rgb folder
    rgb_video_path = "/home/liuyun/HHO-dataset/data_processing/exp_data/20230522_data/24/rgbd/_d435i_camera2_color_image_raw_compressed.mp4"
    # depth folder
    depth_video_path = None
    # time alignment file
    frame_mocap2vision_path = "/home/liuyun/HHO-dataset/data_processing/exp_data/20230522_data/24/frame_mocap2vision.txt"
    # save path
    save_path = "/home/liuyun/HHO-dataset/data_processing/exp_data/20230522_data/24/ex_3rdview_HOI.mp4"
    ############## CHANGE ABOVE #################


    # get data
    mocap_data = get_mocap_data(cfg=cfg, egocentric=False)
    camera_intrinsic, _ = txt2intrinsic(cfg["3rd2_intrinsic_path"])
    camera_pose = np.loadtxt(cfg["3rd2_camera2world_path"])
    mocap_data["camera_pose"] = camera_pose[None, :].repeat(mocap_data["joint_datas"][0].shape[0], axis=0)
    objmodels = [o3d.io.read_triangle_mesh(cfg["obj_model_path"])]
    
    timestamp = []
    with open(rgb_video_path.replace(".mp4", "_timestamp.txt"), "r") as f:
        for line in f:
            values = line.strip().split(",")
            timestamp = [float(v) / 1e9 for v in values]
            break
    cap = cv2.VideoCapture(rgb_video_path)
    fourcc = cv2.VideoWriter_fourcc(*"mp4v")
    fps = cap.get(cv2.CAP_PROP_FPS)
    rgb_data = []
    suc = cap.isOpened()
    frame_cnt = -1
    while True:
        frame_cnt += 1
        suc, img = cap.read()
        if not suc:
            break
        img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
        rgb_data.append((timestamp[frame_cnt], img))
    cap.release()

    # 标注: rgb的第0帧对应mocap数据的第几帧？
    start_mocap_idx = int(np.loadtxt(frame_mocap2vision_path).item())

    # render
    render_mocap_data_on_camera_space(mocap_data["joint_datas"], mocap_data["objposes"], objmodels, mocap_data["camera_pose"], camera_intrinsic, rgb_data=rgb_data, start_mocap_idx=start_mocap_idx, save_path=save_path, fps=30)
